{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "957c80da",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/clutter.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330c50b5",
   "metadata": {
    "id": "A4QOaw_zYLfI",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from IPython.display import clear_output\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    ContactModel,\n",
    "    ContactResults,\n",
    "    ContactVisualizer,\n",
    "    ContactVisualizerParams,\n",
    "    DiagramBuilder,\n",
    "    FixedOffsetFrame,\n",
    "    JointSliders,\n",
    "    LeafSystem,\n",
    "    MeshcatCone,\n",
    "    MeshcatVisualizer,\n",
    "    PointCloud,\n",
    "    PrismaticJoint,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    SpatialInertia,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    "    UnitInertia,\n",
    "    VectorToSkewSymmetric,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import AddShape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ffa2f48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d55f3efc",
   "metadata": {
    "id": "2ctYY0SHJlGB"
   },
   "source": [
    "# Contact wrench cone \"inspector\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "127e5222",
   "metadata": {
    "id": "648WXpNpJlGF"
   },
   "outputs": [],
   "source": [
    "mu = 1.0\n",
    "height = 0.3\n",
    "N = 50\n",
    "t = np.linspace(0, 2 * np.pi, N)\n",
    "vertices = np.vstack(\n",
    "    (height * mu * np.sin(t), height * mu * np.cos(t), height + 0 * t)\n",
    ")\n",
    "vertices = np.append(np.array([[0], [0], [height]]), vertices, axis=1)\n",
    "vertices = np.append(np.zeros((3, 1)), vertices, axis=1)\n",
    "faces = []\n",
    "for i in range(N - 1):\n",
    "    faces.append([0, i + 2, i + 3])\n",
    "    faces.append([1, i + 3, i + 2])\n",
    "faces = np.asarray(faces, dtype=int).T\n",
    "# color = np.tile(np.array([0, 0, 255]), (vertices.shape[1],1)).T\n",
    "cloud = PointCloud(vertices.shape[1])\n",
    "\n",
    "\n",
    "class DrawContactWrench(LeafSystem):\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"contact_results\", AbstractValue.Make(ContactResults())\n",
    "        )\n",
    "        self.DeclareForcedPublishEvent(self.Publish)\n",
    "\n",
    "    def Publish(self, context):\n",
    "        results = self.get_input_port().Eval(context)\n",
    "\n",
    "        for i in range(results.num_point_pair_contacts()):\n",
    "            info = results.point_pair_contact_info(i)\n",
    "            meshcat.SetObject(\n",
    "                f\"contact_{i}\",\n",
    "                MeshcatCone(height, height * mu, height * mu),\n",
    "                rgba=Rgba(0.1, 0.9, 0.1, 1.0),\n",
    "            )\n",
    "            p_WC = info.contact_point()\n",
    "            R_WC = RotationMatrix.MakeFromOneVector(\n",
    "                info.point_pair().nhat_BA_W, 2\n",
    "            )\n",
    "            X_WC = RigidTransform(R_WC, p_WC)\n",
    "            meshcat.SetTransform(f\"contact_{i}\", X_WC)\n",
    "\n",
    "            X_WB = RigidTransform()  # box center is at world origin\n",
    "            meshcat.SetObject(\n",
    "                f\"box_center/contact_{i}/translational\",\n",
    "                MeshcatCone(height, height * mu, height * mu),\n",
    "                rgba=Rgba(0.7, 0.1, 0.1, 1.0),\n",
    "            )\n",
    "            meshcat.SetTransform(\"box_center\", X_WB)\n",
    "            R_BC = X_WB.inverse().rotation() @ R_WC\n",
    "            meshcat.SetTransform(\n",
    "                f\"box_center/contact_{i}\", RigidTransform(R_BC)\n",
    "            )\n",
    "            p_CB_C = -(R_WC @ (X_WB.inverse() @ p_WC))\n",
    "            A_CCtau = VectorToSkewSymmetric(p_CB_C)\n",
    "            # Three.js does not fully support non-uniform scaling\n",
    "            # https://github.com/mrdoob/three.js/issues/15079.\n",
    "            # So I cannot simply draw the meshcatcone and transform it.\n",
    "            # Instead I will create the vertices myself.\n",
    "            cloud.mutable_xyzs()[:] = A_CCtau @ vertices\n",
    "            meshcat.SetObject(\n",
    "                f\"box_center/contact_{i}/rotational_cloud\",\n",
    "                cloud,\n",
    "                point_size=0.002,\n",
    "                rgba=Rgba(0, 1, 0, 1),\n",
    "            )\n",
    "            # TODO(russt): Figure out colors.  The surface mesh example works ok in meshcat_manual_test...\n",
    "            meshcat.SetTriangleMesh(\n",
    "                f\"box_center/contact_{i}/rotational\",\n",
    "                A_CCtau @ vertices,\n",
    "                faces,\n",
    "                rgba=Rgba(0, 0, 1, 1),\n",
    "            )\n",
    "\n",
    "        clear_output(wait=True)\n",
    "\n",
    "\n",
    "def contact_wrench_inspector(second_finger=False):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    box_instance = AddShape(\n",
    "        plant, Box(1, 2, 3), \"box\", color=[0.9, 0.7, 0.5, 0.5]\n",
    "    )\n",
    "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"box\"))\n",
    "    # Specifically opt-in to point contact:\n",
    "    plant.set_contact_model(ContactModel.kPoint)\n",
    "\n",
    "    finger1 = AddShape(\n",
    "        plant, Sphere(0.1), \"finger1\", color=[0.2, 0.2, 0.2, 1.0]\n",
    "    )\n",
    "    box_negative_x = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"box_negative_x\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform([-0.58, 0, 0]),\n",
    "            box_instance,\n",
    "        )\n",
    "    )\n",
    "    finger1_false_body = plant.AddRigidBody(\n",
    "        \"false_body\",\n",
    "        finger1,\n",
    "        SpatialInertia(0, [0, 0, 0], UnitInertia(0, 0, 0)),\n",
    "    )\n",
    "    finger1_y = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"finger1_y\",\n",
    "            box_negative_x,\n",
    "            plant.GetFrameByName(\"false_body\", finger1),\n",
    "            [0, 1, 0],\n",
    "            -0.5,\n",
    "            0.5,\n",
    "        )\n",
    "    )\n",
    "    finger1_z = plant.AddJoint(\n",
    "        PrismaticJoint(\n",
    "            \"finger1_z\",\n",
    "            plant.GetFrameByName(\"false_body\", finger1),\n",
    "            plant.GetFrameByName(\"finger1\"),\n",
    "            [0, 0, 1],\n",
    "            -1.5,\n",
    "            1.5,\n",
    "        )\n",
    "    )\n",
    "    decrement_keycodes = [\"ArrowLeft\", \"ArrowDown\"]\n",
    "    increment_keycodes = [\"ArrowRight\", \"ArrowUp\"]\n",
    "\n",
    "    # optionally add a second finger with the same joints, but on a different face\n",
    "    if second_finger:\n",
    "        finger2 = AddShape(\n",
    "            plant, Sphere(0.1), \"finger2\", color=[0.2, 0.2, 0.2, 1.0]\n",
    "        )\n",
    "        box_positive_x = plant.AddFrame(\n",
    "            FixedOffsetFrame(\n",
    "                \"box_positive_x\",\n",
    "                plant.world_frame(),\n",
    "                RigidTransform([0.58, 0, 0]),\n",
    "                box_instance,\n",
    "            )\n",
    "        )\n",
    "        finger2_false_body = plant.AddRigidBody(\n",
    "            \"false_body\",\n",
    "            finger2,\n",
    "            SpatialInertia(0, [0, 0, 0], UnitInertia(0, 0, 0)),\n",
    "        )\n",
    "        finger2_y = plant.AddJoint(\n",
    "            PrismaticJoint(\n",
    "                \"finger2_y\",\n",
    "                box_positive_x,\n",
    "                plant.GetFrameByName(\"false_body\", finger2),\n",
    "                [0, 1, 0],\n",
    "                -0.5,\n",
    "                0.5,\n",
    "            )\n",
    "        )\n",
    "        finger2_z = plant.AddJoint(\n",
    "            PrismaticJoint(\n",
    "                \"finger2_z\",\n",
    "                plant.GetFrameByName(\"false_body\", finger2),\n",
    "                plant.GetFrameByName(\"finger2\"),\n",
    "                [0, 0, 1],\n",
    "                -1.5,\n",
    "                1.5,\n",
    "            )\n",
    "        )\n",
    "        decrement_keycodes.append(\"KeyA\")\n",
    "        decrement_keycodes.append(\"KeyS\")\n",
    "        increment_keycodes.append(\"KeyD\")\n",
    "        increment_keycodes.append(\"KeyW\")\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "    if False:\n",
    "        cparams = ContactVisualizerParams()\n",
    "        # cparams.force_threshold = 1e-6\n",
    "        cparams.newtons_per_meter = 200.0\n",
    "        cparams.radius = 0.02\n",
    "        contact_visualizer = ContactVisualizer.AddToBuilder(\n",
    "            builder, plant, meshcat, cparams\n",
    "        )\n",
    "\n",
    "    draw_contact_wrench = builder.AddSystem(DrawContactWrench())\n",
    "    builder.Connect(\n",
    "        plant.get_contact_results_output_port(),\n",
    "        draw_contact_wrench.get_input_port(),\n",
    "    )\n",
    "\n",
    "    default_interactive_timeout = None if running_as_notebook else 1.0\n",
    "    sliders = builder.AddSystem(\n",
    "        JointSliders(\n",
    "            meshcat,\n",
    "            plant,\n",
    "            step=[0.05] * (4 if second_finger else 2),\n",
    "            decrement_keycodes=decrement_keycodes,\n",
    "            increment_keycodes=increment_keycodes,\n",
    "        )\n",
    "    )\n",
    "    diagram = builder.Build()\n",
    "    sliders.Run(diagram, default_interactive_timeout)\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "\n",
    "contact_wrench_inspector(second_finger=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f00ac56",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
